import os
import argparse
import pandas as pd
from simnibs import sim_struct, run_simnibs
from datetime import datetime

def tdcs_run_simulation(subpath, pathfem=None, input_dicts=None, csv_path=None):
    """
    Runs SimNIBS simulations using the provided electrode settings or optimized inputs.

    Parameters:
    - subpath: str, required
        Path to the subject's folder.
    - pathfem: str, required
        Path where the simulation outputs will be saved.
    - input_dicts: dict, optional
        Dictionary containing electrode names and their corresponding currents.
        Example: {'Cz': 1.0, 'F3': -1.0}
    - csv_path: str, optional
        Path to a CSV file containing optimized electrode names and currents.
        If provided, this will override `input_dicts`.
    """
    if pathfem==None:
        pathfem=subpath+f"/simulation_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        print("Not assigning pathfem, using",pathfem)
    # Initialize the SimNIBS session
    S = sim_struct.SESSION()
    S.subpath = subpath
    S.pathfem = pathfem
    S.map_to_fsavg = True
    S.map_to_MNI = True
    S.map_to_vol = True
    S.map_to_surf = True
    S.tissues_in_niftis=[1,2]
    
    # Add TDCS list to the session
    tdcs = S.add_tdcslist()

    # Determine electrode settings
    if csv_path is not None:
        print("Using optimized results from CSV instead of input settings.")
        df = pd.read_csv(csv_path, header=None)
        df_short = df[df[1] != 0]
        names = list(df_short[0].values)
        currents = list(df_short[1].values)
    elif input_dicts is not None:
        print("Using input electrode settings.")
        names = list(input_dicts.keys())
        currents = list(input_dicts.values())
    else:
        raise ValueError("Must provide either csv_path or input_dicts,(e.g., '{\"Cz\": 0.001, \"F3\": -0.001}')")

    print(f"Electrode names: {names}")
    print(f"Currents: {currents}")

    # Add electrodes to the simulation
    elec_list = []
    for i in range(len(names)):
        elec_list.append(tdcs.add_electrode())
        elec_list[i].channelnr = i + 1
        elec_list[i].centre = names[i]
        elec_list[i].shape = "ellipse"
        elec_list[i].dimensions = [10, 10]
        elec_list[i].thickness = 4
    
    # Set currents
    tdcs.currents = currents

    # Run the simulation
    print("Running SimNIBS simulation...")
    run_simnibs(S)
    print("Simulation completed.")

if __name__ == "__main__":
    # Set up argument parsing
    parser = argparse.ArgumentParser(description="Run SimNIBS simulations with specified electrode settings.")
    
    parser.add_argument("--subpath", type=str, required=True, help="Path to the subject's folder.")
    parser.add_argument("--pathfem", type=str, help="Path where the simulation outputs will be saved.")
    
    # Optional arguments
    parser.add_argument("--input_dicts", type=str, help="Dictionary containing electrode names and currents (e.g., '{\"Cz\": 0.001, \"F3\": -0.001}')")
    parser.add_argument("--csv_path", type=str, help="Path to a CSV file with optimized electrode settings.")
    
    # Parse the arguments
    args = parser.parse_args()
    
    # Convert input_dicts from string to dict if provided
    input_dicts = None
    if args.input_dicts:
        input_dicts = eval(args.input_dicts)
    
    # Run the simulation with the provided arguments
    tdcs_run_simulation(args.subpath, args.pathfem, input_dicts=input_dicts, csv_path=args.csv_path)
